package com.wgz.mcp.excutor;

import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel;
import com.wgz.mcp.utils.QdrantDocumentUtils;
import io.metaloom.qdrant.client.grpc.QDrantGRPCClient;
import io.metaloom.qdrant.client.grpc.method.GrpcClientRequest;
import io.metaloom.qdrant.client.grpc.proto.JsonWithInt;
import io.metaloom.qdrant.client.grpc.proto.Points;
import io.reactivex.rxjava3.core.Maybe;
import org.springframework.ai.document.Document;

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static io.metaloom.qdrant.client.util.ModelHelper.value;

public class QdRantExcutor {

    public static String queryFromQdrant(
            String tenantId,
            DashScopeEmbeddingModel dashScopeEmbeddingModel,
            String question,
            String qdrantCollectionName,
            String temperature,
            QDrantGRPCClient qDrantGRPCClient) {
        float[] queryEmbedding = null;
        try {
            queryEmbedding = dashScopeEmbeddingModel.embed(question);
        } catch (Exception e) {
            System.out.println(e.getMessage());
        }
        List<Float> queryEmbeddingList = new ArrayList<>();
        for (int i = 0; i < queryEmbedding.length; i++) {
            queryEmbeddingList.add(queryEmbedding[i]);
        }
        // 2. 查询向量数据库

        //配置过滤条件
        Points.Condition authorCondition = Points.Condition.newBuilder()
                .setField(Points.FieldCondition.newBuilder()
                        .setKey("tenant_id")
                        .setMatch(Points.Match.newBuilder()
                                .setKeyword(tenantId.toString())
                                .build())
                        .build())
                .build();

        //根据条件配置过滤器
        Points.Filter filter = Points.Filter.newBuilder()
                .addMust(authorCondition)
                .build();

        //在过滤后的情况下，根据用户问题进行查询
        // 修复：明确指定返回载荷
        Points.WithPayloadSelector payloadSelector = Points.WithPayloadSelector.newBuilder()
                .setEnable(true)  // 启用返回所有载荷字段
                .build();

        // 明确指定不返回向量（节省带宽）
        Points.WithVectorsSelector vectorsSelector = Points.WithVectorsSelector.newBuilder()
                .setEnable(false)  // 不返回向量数据
                .build();

        // 执行查询
        GrpcClientRequest<Points.SearchResponse> searchResponseGrpcClientRequest = qDrantGRPCClient.searchPoints(
                qdrantCollectionName,
                null,                    // vectorName (使用默认)
                queryEmbedding,          // 查询向量
                filter,                  // 租户过滤条件
                Points.SearchParams.newBuilder().build(), // 搜索参数
                1000,                    // limit
                0L,                      // offset
                payloadSelector,         // 修复：返回所有载荷
                vectorsSelector,         // 修复：不返回向量
                Float.valueOf(temperature) // scoreThreshold
        );
        Maybe<Points.SearchResponse> rx = searchResponseGrpcClientRequest.rx();
        Points.SearchResponse searchResponse = rx.blockingGet();
        Stream<Points.ScoredPoint> stream = searchResponse.getResultList().stream();

        Stream<Document> textSegment = stream.map(point -> QdrantDocumentUtils.toDocument(point, "text_segment"));
        List<Document> collect = textSegment.collect(Collectors.toList());

        String context = collect.stream()
                .map(Document::getText)
                .filter(Objects::nonNull)
                .collect(Collectors.joining("\n\n---\n\n"));
        return context;
    }


    public void insertOrUpdateQdrant(
            String vectorId,
            String qdrantCollectionName,
            String tenantId,
            String documentStr,
            QDrantGRPCClient qdrantClient,
            DashScopeEmbeddingModel dashScopeEmbeddingModel
    ) {
        // 1. 使用嵌入模型将文本转换为向量（Embedding）
        float[] embed = dashScopeEmbeddingModel.embed(documentStr);
        List<Float> vector = new ArrayList<>();
        for (int i = 0; i < embed.length; i++) {
            vector.add(embed[i]);
        }
        // 4. 构建Qdrant的数据点（PointStruct）
        Points.PointId pointId = Points.PointId.newBuilder().setUuid(vectorId).build();
//        Points.Vectors.newBuilder().set
        Points.Vectors.Builder vectors = Points.Vectors.newBuilder().setVector(Points.Vector.newBuilder().addAllData(vector).build());

        Points.PointStruct point = Points.PointStruct.newBuilder()
                .setId(pointId) // 设置此点的唯一ID
                .setVectors(vectors) // 设置多维向量
                .putAllPayload( // 设置此点的元数据（负载），是一个或多个键值对 Map
                        Map.of("tenant_id", value(tenantId),
                                "text_segment", value(documentStr),
                                "time", value(new Date().toString())
                        )
                )
                .build();// 构建一个点

        qdrantClient.upsertPoint(qdrantCollectionName, point, false);
    }

    public static Map<String, JsonWithInt.Value> getPointById(String pointId, QDrantGRPCClient qDrantGRPCClient) {
        Points.PointId id = Points.PointId.newBuilder().setUuid(pointId).build();
        GrpcClientRequest<Points.GetResponse> grpcClientRequest = qDrantGRPCClient.getPoint("qdrant_collection_rag", true, true, id);
        Points.GetResponse getResponse = grpcClientRequest.sync();
        if (getResponse != null && getResponse.getResultCount() > 0) {
            Points.RetrievedPoint point = getResponse.getResult(0);
            Map<String, JsonWithInt.Value> payloadMap = point.getPayloadMap();
            return payloadMap;
        } else {
            throw new RuntimeException("未找到ID为 " + pointId + " 的数据");
        }
    }

    public static void writeToQdrant(
            QDrantGRPCClient qDrantGRPCClient
            , DashScopeEmbeddingModel dashScopeEmbeddingModel
            , String text
            , String tenantId
            , String vectorId
    ) {

        float[] embed = dashScopeEmbeddingModel.embed(text);
        ArrayList<Float> list = new ArrayList<>();
        for (int i = 0; i < embed.length; i++) {
            list.add(embed[i]);
        }
        Points.PointStruct point = Points.PointStruct.newBuilder()
                .putAllPayload( // 设置此点的元数据（负载），是一个或多个键值对 Map
                        Map.of("tenant_id", value(tenantId.toString()),
                                "text_segment", value(text),
                                "time", value(new Date().toString())
                        )
                )
                .setVectors(Points.Vectors.newBuilder().setVector(
                        Points.Vector.newBuilder().addAllData(list).build()
                ).build())
                .setId(Points.PointId.newBuilder().setUuid(vectorId).build()).build();
        GrpcClientRequest<Points.PointsOperationResponse> pointsOperationResponseGrpcClientRequest = qDrantGRPCClient.upsertPoint("qdrant_collection_rag", point, false);
        pointsOperationResponseGrpcClientRequest.async();
    }
}
